/**
 * @file h264_to_mp4.c
 * @brief Convert H.264 raw stream to MP4 container format
 * 
 * This tool converts H.264 raw bitstream files to MP4 container format
 * using FFmpeg libraries (libavformat, libavcodec).
 * 
 * @author Your Name
 * @date 2024-03-21
 * @version 1.0
 * 
 * @copyright Copyright (c) 2024
 * 
 * Compilation:
 * gcc -o h264_to_mp4 h264_to_mp4.c -lavformat -lavcodec -lavutil
 * 
 * Usage:
 * ./h264_to_mp4 <input.h264> <output.mp4>
 */

#include <libavformat/avformat.h>
#include <libavcodec/avcodec.h>
#include <libavutil/log.h>
#include <libavutil/timestamp.h>

#define STREAM_FRAME_RATE 25 /* 25 images/s */
#define STREAM_PIX_FMT    AV_PIX_FMT_YUV420P /* default pix_fmt */
/*
typedef struct {
    const char *title;
    const char *author;
    const char *copyright;
    const char *comment;
    const char *creation_time;
} MP4Metadata;
*/
typedef struct {
    const char *title;
    const char *author;
    const char *copyright;
    const char *comment;
    const char *creation_time;
    // 添加新的字段
    const char *album;
    const char *genre;
    const char *year;
    const char *description;
    const char *language;
    const char *rating;
    const char *encoder;
    const char *encoder_version;
    const char *format;
    const char *artist;
    const char *composer;
    const char *publisher;
    const char *producer;
    const char *director;
    const char *keywords;
    const char *location;
    const char *source;
    const char *episode_id;
    const char *network;
    const char *show;
    const char *season;
    const char *episode;
} MP4Metadata;


static void log_packet(const AVFormatContext *fmt_ctx, const AVPacket *pkt)
{
    AVRational *time_base = &fmt_ctx->streams[pkt->stream_index]->time_base;
    av_log(NULL, AV_LOG_DEBUG, "pts:%s pts_time:%s dts:%s dts_time:%s duration:%s duration_time:%s stream_index:%d\n",
           av_ts2str(pkt->pts), av_ts2timestr(pkt->pts, time_base),
           av_ts2str(pkt->dts), av_ts2timestr(pkt->dts, time_base),
           av_ts2str(pkt->duration), av_ts2timestr(pkt->duration, time_base),
           pkt->stream_index);
}

int h264_to_mp4(const char *in_filename, const char *out_filename, const MP4Metadata *metadata)
{
    AVFormatContext *input_ctx = NULL;
    AVFormatContext *output_ctx = NULL;
    AVStream *out_stream = NULL;
    AVCodecParameters *in_codecpar = NULL;
    int ret;
    //int stream_index = 0;
    AVPacket pkt;
    int64_t frame_count = 0;

    // 在打开输入文件之前，设置更大的探测尺寸
    AVDictionary *opts = NULL;
    av_dict_set(&opts, "probesize", "100M", 0);  // 增加到100M
    av_dict_set(&opts, "analyzeduration", "100M", 0);

    // 打开输入文件
    input_ctx = avformat_alloc_context();
    if (!input_ctx) {
        av_log(NULL, AV_LOG_ERROR, "Could not allocate input context\n");
        return AVERROR(ENOMEM);
    }

    // 设置H264解析器
    AVInputFormat *ifmt = av_find_input_format("h264");
    if (!ifmt) {
        av_log(NULL, AV_LOG_ERROR, "Could not find H.264 demuxer\n");
        ret = AVERROR_DEMUXER_NOT_FOUND;
        goto end;
    }

    // 使用选项打开输入
    ret = avformat_open_input(&input_ctx, in_filename, ifmt, &opts);
    if (ret < 0) {
        av_log(NULL, AV_LOG_ERROR, "Could not open input file '%s'\n", in_filename);
        goto end;
    }

    ret = avformat_find_stream_info(input_ctx, NULL);
    if (ret < 0) {
        av_log(NULL, AV_LOG_ERROR, "Failed to retrieve input stream information\n");
        goto end;
    }

    // 创建输出上下文
    avformat_alloc_output_context2(&output_ctx, NULL, NULL, out_filename);
    if (!output_ctx) {
        av_log(NULL, AV_LOG_ERROR, "Could not create output context\n");
        ret = AVERROR_UNKNOWN;
        goto end;
    }

    // 添加视频流
    out_stream = avformat_new_stream(output_ctx, NULL);
    if (!out_stream) {
        av_log(NULL, AV_LOG_ERROR, "Failed allocating output stream\n");
        ret = AVERROR_UNKNOWN;
        goto end;
    }

    in_codecpar = input_ctx->streams[0]->codecpar;

    // 设置输出流参数
    ret = avcodec_parameters_copy(out_stream->codecpar, in_codecpar);
    if (ret < 0) {
        av_log(NULL, AV_LOG_ERROR, "Failed to copy codec params\n");
        goto end;
    }

    // 设置时间基准
    out_stream->time_base = (AVRational){1, STREAM_FRAME_RATE};
    
    if (!(output_ctx->oformat->flags & AVFMT_NOFILE)) {
        ret = avio_open(&output_ctx->pb, out_filename, AVIO_FLAG_WRITE);
        if (ret < 0) {
            av_log(NULL, AV_LOG_ERROR, "Could not open output file '%s'\n", out_filename);
            goto end;
        }
    }

    // 设置元数据
    if (metadata) {
        if (metadata->title)
            av_dict_set(&output_ctx->metadata, "title", metadata->title, 0);
        if (metadata->author)
            av_dict_set(&output_ctx->metadata, "author", metadata->author, 0);
        if (metadata->copyright)
            av_dict_set(&output_ctx->metadata, "copyright", metadata->copyright, 0);
        if (metadata->comment)
            av_dict_set(&output_ctx->metadata, "comment", metadata->comment, 0);
        if (metadata->creation_time)
            av_dict_set(&output_ctx->metadata, "creation_time", metadata->creation_time, 0);

        if (metadata->album)
            av_dict_set(&output_ctx->metadata, "album", metadata->album, 0);
        if (metadata->genre)
            av_dict_set(&output_ctx->metadata, "genre", metadata->genre, 0);
        if (metadata->year)
            av_dict_set(&output_ctx->metadata, "year", metadata->year, 0);
        if (metadata->description)
            av_dict_set(&output_ctx->metadata, "description", metadata->description, 0);
        if (metadata->language)
            av_dict_set(&output_ctx->metadata, "language", metadata->language, 0);

        if (metadata->rating)
            av_dict_set(&output_ctx->metadata, "rating", metadata->rating, 0);
        if (metadata->encoder)
            av_dict_set(&output_ctx->metadata, "encoder", metadata->encoder, 0);
        if (metadata->encoder_version)
            av_dict_set(&output_ctx->metadata, "encoder_version", metadata->encoder_version, 0);
        if (metadata->format)
            av_dict_set(&output_ctx->metadata, "format", metadata->format, 0);

        if (metadata->artist)
            av_dict_set(&output_ctx->metadata, "artist", metadata->artist, 0);
        if (metadata->composer)
            av_dict_set(&output_ctx->metadata, "composer", metadata->composer, 0);
        if (metadata->publisher)
            av_dict_set(&output_ctx->metadata, "publisher", metadata->publisher, 0);
        if (metadata->producer)
            av_dict_set(&output_ctx->metadata, "producer", metadata->producer, 0);
        if (metadata->director)
            av_dict_set(&output_ctx->metadata, "director", metadata->director, 0);

        if (metadata->keywords)
            av_dict_set(&output_ctx->metadata, "keywords", metadata->keywords, 0);
        if (metadata->location)
            av_dict_set(&output_ctx->metadata, "location", metadata->location, 0);
        if (metadata->source)
            av_dict_set(&output_ctx->metadata, "source", metadata->source, 0);
        if (metadata->episode_id)
            av_dict_set(&output_ctx->metadata, "episode_id", metadata->episode_id, 0);
        if (metadata->network)
            av_dict_set(&output_ctx->metadata, "network", metadata->network, 0);
        if (metadata->show)
            av_dict_set(&output_ctx->metadata, "show", metadata->show, 0);
        if (metadata->season)
            av_dict_set(&output_ctx->metadata, "season", metadata->season, 0);
        if (metadata->episode)
            av_dict_set(&output_ctx->metadata, "episode", metadata->episode, 0);
    }

    // 写入文件头
    ret = avformat_write_header(output_ctx, NULL);
    if (ret < 0) {
        av_log(NULL, AV_LOG_ERROR, "Error occurred when writing header\n");
        goto end;
    }

    // 读取所有数据包
    while (1) {
        ret = av_read_frame(input_ctx, &pkt);
        if (ret < 0)
            break;

        pkt.stream_index = 0;
        
        // 设置正确的时间戳
        if (pkt.pts == AV_NOPTS_VALUE) {
            pkt.pts = frame_count;
            pkt.dts = frame_count;
            frame_count++;
        }

        // 时间戳转换
        pkt.pts = av_rescale_q_rnd(pkt.pts,
                                  (AVRational){1, STREAM_FRAME_RATE}, // 使用固定帧率
                                  out_stream->time_base,
                                  AV_ROUND_NEAR_INF|AV_ROUND_PASS_MINMAX);
        pkt.dts = pkt.pts;  // 对于H264，通常dts = pts
        
        // 设置持续时间
        pkt.duration = av_rescale_q(1,
                                  (AVRational){1, STREAM_FRAME_RATE},
                                  out_stream->time_base);
        pkt.pos = -1;

        log_packet(output_ctx, &pkt);

        ret = av_interleaved_write_frame(output_ctx, &pkt);
        if (ret < 0) {
            char errbuf[AV_ERROR_MAX_STRING_SIZE];
            av_strerror(ret, errbuf, AV_ERROR_MAX_STRING_SIZE);
            av_log(NULL, AV_LOG_ERROR, "Error muxing packet: %s\n", errbuf);
            break;
        }
        av_packet_unref(&pkt);
    }

    // 写入文件尾
    av_write_trailer(output_ctx);

end:
    if (opts)
        av_dict_free(&opts);
    if (input_ctx)
        avformat_close_input(&input_ctx);

    if (output_ctx && !(output_ctx->oformat->flags & AVFMT_NOFILE))
        avio_closep(&output_ctx->pb);

    if (output_ctx)
        avformat_free_context(output_ctx);

    return ret < 0 ? 1 : 0;
}

int main(int argc, char **argv)
{
    if (argc != 3) {
        fprintf(stderr, "Usage: %s <input h264 file> <output mp4 file>\n", argv[0]);
        return 1;
    }
    /*
    MP4Metadata metadata = {
        .title = "My Video",
        .author = "Your Name",
        .copyright = "Copyright 2024",
        .comment = "Created with FFmpeg",
        .creation_time = "2024-03-21"
    };
    */
    MP4Metadata metadata = {
        // 基本信息
        .title = "示例视频",
        .author = "张三",
        .copyright = "版权所有 2024",
        .comment = "使用FFmpeg创建",
        .creation_time = "2024-03-21T10:00:00Z",
        
        // 媒体信息
        .album = "视频合集2024",
        .genre = "纪录片",
        .year = "2024",
        .description = "这是一个示例视频，展示FFmpeg的使用",
        .language = "chi",
        
        // 分级和技术信息
        .rating = "PG",
        .encoder = "FFmpeg",
        .encoder_version = "4.2.7",
        .format = "H.264/MP4",
        
        // 创作者信息
        .artist = "视频创作组",
        .composer = "李四",
        .publisher = "XX工作室",
        .producer = "王五",
        .director = "赵六",
        
        // 其他元数据
        .keywords = "示例,教程,FFmpeg",
        .location = "北京",
        .source = "原创",
        .episode_id = "EP001",
        .network = "测试网络",
        .show = "FFmpeg教程",
        .season = "S01",
        .episode = "E01"
    };

    av_log_set_level(AV_LOG_ERROR);
    return h264_to_mp4(argv[1], argv[2], &metadata);
} 


/*
编译选项
gcc -o h264_to_mp4 h264_to_mp4.c -lavformat -lavcodec -lavutil


查看元数据的方法

# 使用ffprobe查看
ffprobe -show_format output.mp4

# 使用mediainfo查看更详细信息
mediainfo output.mp4

# 使用exiftool查看
exiftool output.mp4

sudo apt install libimage-exiftool-perl

*/
